import torch
import uvicorn
from asgiref.sync import sync_to_async
from fastapi import FastAPI, Form, status, HTTPException, UploadFile, File
from fastapi.responses import JSONResponse
import logging
import os

from pydantic import BaseModel

from indextts.infer_spk import IndexTTS

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
app = FastAPI(title="TTS API")

tts = IndexTTS(model_dir="checkpoints",cfg_path="checkpoints/config.yaml")

class TTS_Request(BaseModel):
    spk: str = None
    tts_text: str = None
    save_dir: str = "audio"
    title: str = None

@app.post("/generateAudio", summary="语音生成")
async def generate(request:TTS_Request):
    req=request.model_dump()
    spk=req.get('spk')
    tts_text=req.get('tts_text')
    output_path=req.get('save_dir')
    title=req.get('title')

    # result=generate_instance.generate_audio(spk=spk,tts_text=text, instruct_text=instruct_text, stream=stream, speed=speed,
    #                           save_dir=save_dir)
    wavs=[]
    for text in tts_text.split('\n'):
        if text.strip()=='':
            continue
        else:
            wav = await sync_to_async(tts.infer)(spk, text, output_path,title)
            wavs.append(wav)
    if len(wavs)==0:
        return JSONResponse(content={"description": "No text", "status": "failed"},status_code=status.HTTP_400_BAD_REQUEST)

    if len(wavs)>0:
        result = torch.cat(wavs, dim=1)
        response={"description": "Generate completed", "status": "success",'wav_path':None}
        if result:
            response['wav_path']=result
            logging.info(f"generate audio save :{result}")
            return JSONResponse(content=response,status_code=status.HTTP_200_OK)
        else:
            response = {"description": f"generate error", "status": "failed"}
            raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=response)
    else:
        response = {"description": f"generate error", "status": "failed"}
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=response)


@app.post("/train", summary="声音保存")
async def train(spk: str = Form(...), audio: UploadFile = File(...)):
    try:
        # 设置保存文件的目录和文件名
        save_dir = "received"  # 保存文件的目录
        os.makedirs(save_dir, exist_ok=True)  # 如果目录不存在，创建它

        audio_prompt = os.path.join(save_dir, audio.filename)  # 保存路径

        # 保存上传的文件
        with open(audio_prompt, "wb") as buffer:
            buffer.write(await audio.read())  # 读取文件内容并写入

        # 调用 tts.extract 处理文件
        rest = tts.extract(audio_prompt, spk)  # 假设 tts.extract 接收文件路径和说话者参数

        if rest:
            return JSONResponse(content={"message": "保存成功"},
                                status_code=status.HTTP_201_CREATED)
        else:
            return JSONResponse(content={"message": "保存失败"},
                                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
    except Exception as e:
        logging.error(f"保存失败: {e}")
        raise HTTPException(status_code=500, detail="保存失败")
# 启动服务
if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=7895)
